#ifndef CDETECTION_HPP
#define CDETECTION_HPP
#include "Hungarian.hpp"
#include "common.h"
#include <iostream>
#include <map>
#include <opencv2/opencv.hpp>
#include <vector>

using namespace std;
using namespace cv;

class CDetection {
private:
  /* data */
  void Track(vector<KeyPoint> &keypts);
  double GetIOU(Rect_<float> &a, Rect_<float> &b);
  bool IsInside(Rect_<float> &a, Rect_<float> &b);
  void GetOutPut(vector<KeyPoint> &keypts, vector<Point2f> &pts);
  bool IsBlob(KeyPoint &key, Mat &img, float res_T);
  Ptr<SimpleBlobDetector> sbd;
  vector<Rect_<float>> trks;
  vector<Rect_<float>> dets;
  vector<KeyPoint> tracked_dets;
  vector<vector<double>> iou_matrix;
  vector<int> assignment;
  set<int> unmatched_dets;
  set<int> unmatched_trks;
  set<int> all_items;
  set<int> matched_items;
  vector<Point> matched_pairs;
  HungarianAlgorithm HungAlgo;

public:
  CDetection(/* args */);
  ~CDetection();
  void Run(Mat &img, vector<Point2f> &pts);
  double iou_T_h, iou_T_l;
  bool mNeedReognition;
};

CDetection::CDetection(/* args */) {
  iou_T_h = 0.2;
  iou_T_l = 0.05;
  SimpleBlobDetector::Params pDefaultBLOB;
  // This is default parameters for SimpleBlobDetector
  pDefaultBLOB.thresholdStep = 20;
  pDefaultBLOB.minThreshold = 100;
  pDefaultBLOB.maxThreshold = 220;
  pDefaultBLOB.minRepeatability = 2;
  pDefaultBLOB.minDistBetweenBlobs = 5;
  pDefaultBLOB.filterByColor = true;
  pDefaultBLOB.blobColor = 255;
  pDefaultBLOB.filterByArea = true;
  pDefaultBLOB.minArea = 30;
  pDefaultBLOB.maxArea = 400;
  pDefaultBLOB.filterByCircularity = false;
  pDefaultBLOB.minCircularity = 0.9f;
  pDefaultBLOB.maxCircularity = (float)1e37;
  pDefaultBLOB.filterByInertia = false;
  pDefaultBLOB.minInertiaRatio = 0.1f;
  pDefaultBLOB.maxInertiaRatio = (float)1e37;
  pDefaultBLOB.filterByConvexity = true;
  pDefaultBLOB.minConvexity = 0.70f;
  pDefaultBLOB.maxConvexity = (float)1e37;
  Ptr<Feature2D> b = SimpleBlobDetector::create(pDefaultBLOB);
  sbd = b.dynamicCast<SimpleBlobDetector>();

  mNeedReognition = true;
}

CDetection::~CDetection() {}

bool CDetection::IsBlob(KeyPoint &key, Mat &img, float res_T) {
  bool is_blob = true;
  int width = img.cols;
  int height = img.rows;
  vector<Point> pts;
  // center
  int center_x = key.pt.x, center_y = key.pt.y;
  pts.push_back(Point(center_x, center_y));
  // center left
  int tmp_x, tmp_y;
  tmp_x = center_x - 2;
  tmp_y = center_y;
  pts.push_back(Point(tmp_x, tmp_y));
  // center top
  tmp_x = center_x;
  tmp_y = center_y - 2;
  pts.push_back(Point(tmp_x, tmp_y));
  // center right
  tmp_x = center_x + 2;
  tmp_y = center_y;
  pts.push_back(Point(tmp_x, tmp_y));
  // center bottom
  tmp_x = center_x;
  tmp_y = center_y + 2;
  pts.push_back(Point(tmp_x, tmp_y));

  // edge top left
  tmp_x = center_x - key.size;
  tmp_y = center_y - key.size;
  pts.push_back(Point(tmp_x, tmp_y));
  // edge top right
  tmp_x = center_x + key.size;
  tmp_y = center_y - key.size;
  pts.push_back(Point(tmp_x, tmp_y));
  // edge bottom riight
  tmp_x = center_x + key.size;
  tmp_y = center_y + key.size;
  pts.push_back(Point(tmp_x, tmp_y));
  // edge bottom left
  tmp_x = center_x - key.size;
  tmp_y = center_y + key.size;
  pts.push_back(Point(tmp_x, tmp_y));

  double sum = 0.0, center_mean = 0.0;
  double value = 0.0;
  // cout << "value: ";
  for (auto i = 0; i < pts.size(); ++i) {
    tmp_x = min(max(0, pts[i].x), width - 1);
    tmp_y = min(max(0, pts[i].y), height - 1);
    pts[i] = Point(tmp_x, tmp_y);
    // value = (double)(img.at<Vec3b>(tmp_y, tmp_x)[0])*0.11
    //         + (double)(img.at<Vec3b>(tmp_y, tmp_x)[1])*0.59
    //         + (double)(img.at<Vec3b>(tmp_y, tmp_x)[2])*0.30;

    value = (double)img.at<uchar>(tmp_y, tmp_x);
    // cout << value << " ";
    if (i <= 4) {
      sum += value;
    }
    if (i == 4)
      center_mean = sum / 5.0;
    if (i > 4) {
      if (center_mean - value < res_T * value) {
        is_blob = false;
        break;
      }
    }
  }
  return is_blob;
}
void CDetection::Run(Mat &img, vector<Point2f> &pts) {
  CTimer detec_time("Detec");
  CTimer select_time("Select");
  CTimer track_time("TRACK");
  CTimer filter_time("FILTER");
  if (TIME)
    filter_time.start();
  GaussianBlur(img, img, Size(5, 5), 0.0);
  if (TIME)
    filter_time.end();
  vector<KeyPoint> key_dets_ori;
  vector<KeyPoint> key_dets;
  vector<Rect> zone;
  vector<vector<Point>> region;
  static int video_count = 0;

  if (!img.empty()) {
    if (TIME)
      detec_time.start();
    sbd->detect(img, key_dets_ori, Mat());
    if (TIME)
      detec_time.end();
    // drawKeypoints(img, key_dets, img);
    if (TIME)
      select_time.start();
    // cout << "ori_size: " << key_dets_ori.size() << endl;
    for (vector<KeyPoint>::iterator k = key_dets_ori.begin();
         k != key_dets_ori.end(); ++k) {
      // if (IsBlob(*k, img, 0.15)) {
        key_dets.push_back(*k);
        // cout << "octave: " << k->octave << ", response: " << k->response << ", size: " << k->size << endl;
        // circle(img, k->pt, (int)k->size, RED, 2);
        // putText(img, to_string(k->octave)+","+to_string(k->response),
        // Point(k->pt.x-10, k->pt.y+20), cv::FONT_HERSHEY_COMPLEX, 0.5, GREEN,
        // 1);
      // }
    }
    if (TIME)
      select_time.end();
    if (TIME)
      track_time.start();
    if (!mNeedReognition) {
      for (auto key : key_dets) {
        pts.push_back(key.pt);
      }
    } else {
      // track
      if (trks.size() > 0) {
        Track(key_dets);
        GetOutPut(tracked_dets, pts);
      } else {
        GetOutPut(key_dets, pts);
      }
      dets.clear();
      for (auto key : key_dets) {
        Rect_<float> det = Rect_<float>(key.pt.x, key.pt.y, key.size, key.size);
        dets.push_back(det);
      }
      trks.swap(dets);
      if (TIME)
        track_time.end();
    }
  }
  video_count++;
}

void CDetection::Track(vector<KeyPoint> &dets) {
  unmatched_trks.clear();
  unmatched_dets.clear();
  all_items.clear();
  matched_items.clear();
  assignment.clear();
  tracked_dets.clear();

  int trk_num = trks.size();
  int det_num = dets.size();
  iou_matrix.clear();
  iou_matrix.resize(trk_num, vector<double>(det_num, 0));
  for (int i = 0; i < trk_num; i++) // compute iou matrix as a distance matrix
  {
    for (int j = 0; j < det_num; j++) {
      Rect_<float> det =
          Rect_<float>(dets[j].pt.x, dets[j].pt.y, dets[j].size, dets[j].size);

      // use 1-iou because the hungarian algorithm computes a minimum-cost
      // assignment.
      iou_matrix[i][j] = 1 - GetIOU(trks[i], det);
    }
  }

  // solve the assignment problem using hungarian algorithm.
  // the resulting assignment is [track(prediction) : detection], with
  // len=preNum
  if (iou_matrix.size() > 0) {
    HungAlgo.Solve(iou_matrix, assignment);
  }
  // find matches, unmatched_detections and unmatched_predictions
  if (det_num > trk_num) // there are unmatched detections
  {
    for (int n = 0; n < det_num; n++)
      all_items.insert(n);
    for (int i = 0; i < trk_num; i++)
      matched_items.insert(assignment[i]);
    set_difference(
        all_items.begin(), all_items.end(), matched_items.begin(),
        matched_items.end(),
        insert_iterator<set<int>>(unmatched_dets, unmatched_dets.begin()));
  } else if (det_num < trk_num) // there are unmatched trajectory/predictions
  {
    for (int i = 0; i < trk_num; i++)
      if (assignment[i] ==
          -1) // unassigned label will be set as -1 in the assignment algorithm
        unmatched_trks.insert(i);
  } else
    ;
  // filter out matched with low IOU
  matched_pairs.clear();
  for (int i = 0; i < trk_num; i++) {
    if (assignment[i] == -1) // pass over invalid values
      continue;
    if (1 - iou_matrix[i][assignment[i]] < iou_T_h) {
      int tmp = assignment[i];
      Rect_<float> det = Rect_<float>(dets[tmp].pt.x, dets[tmp].pt.y,
                                      dets[tmp].size, dets[tmp].size);

      if (!IsInside(trks[i], det)) {
        if (1 - iou_matrix[i][assignment[i]] < iou_T_l) {
          unmatched_trks.insert(i);
          unmatched_dets.insert(assignment[i]);
        }
      }
    } else {
      matched_pairs.push_back(cv::Point(i, assignment[i]));
    }
  }
  if (unmatched_dets.size() > 0) {
    for (auto i : unmatched_dets) {
      tracked_dets.push_back(dets[i]);
    }
  }
}

double CDetection::GetIOU(Rect_<float> &a, Rect_<float> &b) {
  float in = (a & b).area();
  float un = a.area() + b.area() - in;
  float iou = 0.0;

  if (un > DBL_EPSILON)
    iou = (in / un);
  float dis = sqrt(pow(a.x - b.x, 2) + pow(a.y - b.y, 2));
  float dis_max = max(a.width, b.width) * 3;
  float dis_w = 1 - dis / dis_max;
  return (iou + dis_w) * 0.5;
}
bool CDetection::IsInside(Rect_<float> &a, Rect_<float> &b) {
  float in = (a & b).area();
  // cout << "in: " << in << " a: " << a.area() << " b: " << b.area() << endl;
  if (abs(in - a.area()) <= DBL_EPSILON || abs(in - b.area()) <= DBL_EPSILON) {
    return true;
  }
  return false;
}

void CDetection::GetOutPut(vector<KeyPoint> &keypts, vector<Point2f> &pts) {
  if (keypts.size() > 0) {
    for (auto pt : keypts) {
      pts.push_back(pt.pt);
    }
  }
}
#endif
